package gbench.sandbox.matlib;

import org.junit.jupiter.api.Test;

import java.util.*;
import java.util.function.Function;
import java.util.function.UnaryOperator;
import java.util.stream.*;
import gbench.common.matlib.MatlibCanvas.NVec;
import gbench.common.matlib.MatlibCanvas.Range;
import gbench.common.matlib.MatlibCanvas.XVec;
import gbench.common.matlib.algebra.op.Ops;
import gbench.common.matlib.matrix.MatrixOps;
import gbench.common.matlib.algebra.lisp.*;

import static gbench.common.matlib.MatlibCanvas.println;
import static gbench.common.matlib.algebra.lisp.MyRecord.*;
import static gbench.common.matlib.algebra.lisp.Tuple2.TUP2;

/**
 * 节点的网格运算
 * 
 * @author gbench
 *
 */
public class JunitCell {

    /**
     * 计算节点
     * 
     * @author gbench
     *
     */
    class Cell {

        /**
         * 节点序号
         * 
         * @param id
         */
        public Cell(int i) {
            this.id = i;
        }

        /**
         * 获取计算节点名称
         * 
         * @return 计算节点名称
         */
        public String getName() {
            return "#" + this.id;
        }

        /**
         * 计算单元初始化
         * 
         * @return Cell 本身
         */
        public Cell initialize() {
            return this.initialize(this.inputs.size());
        }

        /**
         * 计算单元初始化
         * 
         * @param size 输入数据的数量
         * @return Cell 本身
         */
        @SuppressWarnings("unchecked")
        public Cell initialize(final int size) {
            if (size < 1) {
                return this;
            }

            final Function<Integer, Double> init = cell -> Optional.ofNullable(this.dblAttr("init")) // 提取初始值
                    .orElse(Math.random()); // 权重初始化
            Stream.iterate(0, i -> i + 1).limit(size).forEach(i -> this.weights.add(init.apply(i)));
            this.bias = init.apply(-1); // bias 标号为-1
            this.activation = (UnaryOperator<Double>) Optional.ofNullable(this.attr("activation")).orElse(this.activation);

            return this;
        }

        public List<Cell> getInputs() {
            return this.inputs;
        }

        public Stream<Cell> getInputS() {
            return this.inputs.stream();
        }

        public Cell addInput(Cell input) {
            this.inputs.add(input);
            return this;
        }

        /**
         * 是否初始化过
         * 
         * @return 是否初始化过
         */
        public boolean ready() {
            return this.bias == null;
        }

        /**
         * 节点的计算函数
         * 
         * @param x 输入的数据向量
         * @return Double
         */
        public Double evaluate(final Number... x) {
            return this.evaluate(NVec.of(x));
        }

        /**
         * 节点的计算函数
         * 
         * @param x 输入的数据向量
         * @return Double
         */
        public Double evaluate(final NVec x) {

            if (this.attr("value") != null) {
                return this.dblAttr("value"); // 使用缓存的计算结果
            }

            if (this.ready()) { // bias 为空时是否初始化的标志
                final var size = this.inputs.size();
                this.initialize(size == 0 ? x.size() : size);
            } // if

            final var _x = this.inputs.size() > 0
                    ? this.getInputS().map(e -> e.evaluate(x)).collect(NVec.nveclc())
                    : x;
            final var z = _x.dot(NVec.of(this.weights)) + this.bias;
            final var a = activation.apply(z); // final 结果

            this.attr("z", z); // 累加值
            this.attr("a", a); // 累加值
            this.attr("value", a); // 设置计算的值

            return a;
        }
        
        /**
         * 层号
         * @return
         */
        public Integer level() {
            return this.intAttr("layer");
        }
        
        /**
         * 设置层号
         * @return
         */
        public Integer level(final int level) {
            this.attr("layer", level);
            return level;
        }

        /**
         * 提取属性
         * 
         * @param name 属性名称
         * @return 属性值
         */
        public Integer intAttr(final String name) {
            return this.attributes.i4(name);
        }

        /**
         * 提取属性
         * 
         * @param name 属性名称
         * @return 属性值
         */
        public Double dblAttr(final String name) {
            return this.attributes.dbl(name);
        }

        /**
         * 提取属性
         * 
         * @param name 属性名称
         * @return 属性值
         */
        public String strAttr(final String name) {
            return this.attributes.str(name);
        }

        /**
         * 添加属性
         * 
         * @param name  属性名
         * @param value 属性值
         * @return 属性集的当前的IRecord 以保证可以链式编程。
         */
        public IRecord attr(final String name, final Object value) {
            return this.attributes.add(name, value);
        }

        /**
         * 提取属性
         * 
         * @param name 属性名称
         * @return 属性值
         */
        public Object attr(final String name) {
            return this.attributes.get(name);
        }

        /**
         * 属性回调
         * 
         * @param <T>    元素类型
         * @param mapper 属性计算的回调函数
         * @return mapper 返回值
         */
        public <T> T attr(final Function<IRecord, T> mapper) {
            return mapper.apply(this.attributes);
        }

        /**
         * 参数更新
         * 
         * @param updator 权重更新函数 (i:参数索引,i==-1 --> bias ,i>=0 --> weights索引 )->参数更新值
         * @return Cell 本身
         */
        public Cell update(final Function<Integer, Double> updator) {
            this.bias = updator.apply(-1); // 更新偏移值
            for (int i = 0; i < this.weights.size(); i++) {
                this.weights.set(i, updator.apply(i)); // 权重更新
            }
            return this;
        }
        
        /**
         * 计算节点的状态数据清空
         * 
         * @return Cell 对象本身
         */
        public Cell clear() {
            this.bias = null;
            if (this.weights != null) {
                this.weights.clear();
            }
            if (this.attributes != null) {
                this.attributes.toMap().clear();
            }
            return this;
        }

        /**
         * 字符串格式化
         */
        public String toString() {
            return "cell " + id;
        }

        private int id; // 结算节点的id
        private Double bias = null; // 偏移量
        private List<Cell> inputs = new ArrayList<Cell>(); // 输入数据的cells
        private List<Double> weights = new ArrayList<Double>(); // 数据权重
        private IRecord attributes = REC(); // 计算节点的属性
        private UnaryOperator<Double> activation = x -> 1 / (1 + Math.exp(-x)); // 激活函数
    }
    
    /**
     * 节点建模
     * 
     * @param cells 计算节点
     * @param nums 各个层次的节点数量
     * @return 分层之后的网络
     */
    public Map<Integer, List<IRecord>> sequenceModel(final XVec<Cell> cells, final int... nums) {
        // 计算网络的设计
        cells.sliceX(nums) // 划分层节点
            .fmap(Ops.kvp_int()).fmap(e -> { // 设置节点的层号信息,(level,layer)
                final var level = e._1; // 层号
                e._2.foreach(cell -> { // 节点的属性信息设置
                    cell.attr("layer", level) // 设置层号
                        .add("init", 1) // 初始值
                        .add("activation", (UnaryOperator<Double>) x -> x); // 激活函数
                }); // foreach
                return e._2(); // 返回节点层
            }).slidingS(2, 1, true) // 提取比邻关系的层[前一层,当前层]
            .map(Range::cph) // 构造比邻层节点的间的全连接 网络
            .forEach(edges -> { // 设置层与层之间的节点关系
                edges.forEach(edge -> { // 设置计算节点的彼此输入关系
                    edge[1].addInput(edge[0]); // 1: 后层节点, 0: 前层节点
                }); // forEach
            }); // forEach

        final var node_rb = IRecord.rb("layer,cell"); // 网络节点 builder
        final var layerMap = cells.map(e -> node_rb.get(e.intAttr("layer"), e))
                .collect(Collectors.groupingBy(e -> e.i4("layer"))); // 网络结构
        return layerMap;
    }

    @Test
    public void foo() {
        final var cells = XVec.OF(10, Cell::new); // 准备计算节点
        final var model = this.sequenceModel(cells, 1, 2, 3, 2, 2); // 建立计算网络

        final var data = NVec.of(10, i -> 1); // 数据输入
        final var output_rb = IRecord.rb("name,layer,value"); // 输出层节点 builder
        final var inputs = cells.flatMap(Cell::getInputS).collect(Collectors.toSet()); // 有输入能力的节点
        final var outputs = cells.filter(e -> !inputs.contains(e)) // 网络的输出节点
                .map(e -> output_rb.get(e.getName(), e.attr("layer"), e.evaluate(data))) // 计算输出节点
                .collect(DFrame.dfmclc);
        println("输出节点：");
        println(outputs);

        final var cell_rb = IRecord.rb("layer,name,value,inputs,bias,weights,attribute"); // cell 节点 builder
        println("网络状态:");
        model.forEach((level, layer) -> {
            final var dfm2 = layer.stream().map(node -> { // 层节点
                final var id = node.i4("layer"); // 层号
                final var cell = node.get("cell", Cell.class); // 提取运算节点
                return cell_rb.get(id, cell.getName(), cell.attr("value"), cell.inputs, cell.bias, cell.weights, cell.attributes);
            }).collect(DFrame.dfmclc);
            println("------------------------------------------\n" + 
                    "#layer:" + level + "\n"
                    + "---------------------------------------");
            println(dfm2);
        }); // forEach
        
        println("=======================================================================================");
        final XVec<XVec<Cell>> seqModel = model.entrySet().stream().map(Map.Entry::getValue)
                .map(layer -> XVec.OF(layer).fmap(e -> e.get("cell", Cell.class))) // 提取计算节点
                .collect(XVec.xveclc()); // 神经网络
        final var y = seqModel // 提取各个层计算节点
                .reduce(data, (acc, layerX) -> { // 分层计算
                    final var activations = new ArrayList<UnaryOperator<Double>>(); // 各个节点的激活函数
                    final var wts = layerX.fmap(cell -> { // 提取计算节点信息
                        activations.add(cell.activation); // 记录activation函数
                        return NVec.of(cell.weights).insert(0, cell.bias).data();
                    }).toArray(); // 生成节点权重矩阵
                    final var z = MatrixOps.mmult(wts, acc.insert(0, 1d).data()); // acc补充bias 后 与权重矩阵相乘 生成z向量
                    
                    return NVec.of(z).fmapN((i, e) -> activations.get(i).apply(e)); // z 向量转换成 新的acc向量
                });
        println(y);
        
        // back pass
        final var label_y = y.duplicate(); // 标准值
        final var grad_cost_of_y= label_y.sub(y).mutateN(e -> e.mul(-1)); // cost 函数
        final var t = seqModel.reverse().slidingS(2, 1, true).map(e -> {
            final var level = e.car().car().level();
            final var pre = e.car();
            final var cur = e.get(1);
            final var wtx = pre.fmap(c -> NVec.of(c.weights).insert(0, c.bias).toArray()).toArray();
            final var zl = cur.fmap(c -> c.dblAttr("z")).fmap(x -> {
                final var sigmoid = 1.0 / (1.0 + Math.exp(-x));
                return sigmoid * (1 - sigmoid);
            }).toArray();
            println(level, zl, pre, cur);
            return TUP2(zl, MatrixOps.transpose(wtx));
        }).reduce(grad_cost_of_y, (acc, c) -> {
            final var d = MatrixOps.mmult(c._1, c._2);
            final var _t = NVec.of(d).mul(acc);
            println();
            return _t;
        }, (a, b) -> a);
        println(t);
        
        println("=======================================================================================");
        // forward pass
        seqModel.reduce(data, (acc, layerX) -> {
            final var activations = new ArrayList<UnaryOperator<Double>>(); // 各个节点的激活函数
            final var wts = layerX.fmap(cell -> { // 提取计算节点信息
                activations.add(x -> x); // 记录activation函数
                return NVec.of(cell.weights).insert(0, cell.bias).data();
            }).toArray(); // 生成节点权重矩阵

            println("--\n",MatrixOps.fmt(wts));
            return acc;
        });
    }

}
